from typing import List, Dict, Any, Optional, Tuple
from datetime import datetime, timedelta
from enum import Enum
from utils.logger import Logger
from utils.exceptions import ToolError, ParamError
from middleware.redis_adapter import get_redis_adapter
from utils.constants import TASK_STATUSES, DEFAULT_DATETIME_FORMAT
from apps.orchestrator.smart_task_distributor import smart_task_distributor

# 初始化日志
logger = Logger.get_logger(agent_name="task_scheduler_agent")


# 任务优先级枚举（与系统常量对齐）
class TaskPriority(Enum):
    HIGH = "high"  # 高优先级（如紧急查询、多模态任务）
    MEDIUM = "medium"  # 中优先级（如文档问答、工具调用）
    NORMAL = "normal"  # 普通优先级（如文本问答）
    LOW = "low"  # 低优先级（如日志分析、统计任务）


class TaskSchedulerAgent:
    """任务调度Agent：基于Redis实现分布式任务队列，支持优先级和超时管理"""

    def __init__(self):
        self.redis_client = get_redis_adapter()  # 使用Redis适配器，支持集群模式
        # Redis键前缀（区分不同队列）
        self.queue_prefix = "task_queue:"
        self.task_meta_prefix = "task_meta:"  # 任务元信息存储键
        self.max_queue_size = 1000  # 最大队列长度（防止内存溢出）
        self.task_timeout = 300  # 任务默认超时时间（5分钟）
        self.use_smart_distribution = True  # 是否使用智能任务分发
        logger.info("TaskSchedulerAgent initialized successfully")

    def _get_queue_key(self, priority: TaskPriority) -> str:
        """获取对应优先级的队列键"""
        return f"{self.queue_prefix}{priority.value}"

    def _validate_task(self, task: Dict[str, Any]) -> None:
        """验证任务格式合法性"""
        required_fields = ["task_id", "task_type", "query"]
        missing_fields = [f for f in required_fields if f not in task]
        if missing_fields:
            raise ParamError(
                message=f"任务缺少必填字段: {missing_fields}",
                context={"task": task}
            )

        # 验证优先级
        priority = task.get("metadata", {}).get("priority", "normal")
        if priority not in [p.value for p in TaskPriority]:
            raise ParamError(
                message=f"无效的任务优先级: {priority}，支持：{[p.value for p in TaskPriority]}",
                context={"task_id": task.get("task_id")}
            )

    def enqueue_task(self, task: Dict[str, Any]) -> bool:
        """
        任务入队（按优先级加入对应队列）
        Args:
            task: 任务字典（需包含task_id、task_type、query，可选metadata）
        Returns:
            入队成功返回True，失败返回False
        """
        try:
            # 1. 任务验证
            self._validate_task(task)
            task_id = task["task_id"]
            metadata = task.get("metadata", {})
            priority = TaskPriority(metadata.get("priority", "normal"))
            queue_key = self._get_queue_key(priority)

            # 2. 检查队列长度（防止溢出）
            queue_size = self.redis_client.llen(queue_key)
            if queue_size >= self.max_queue_size:
                logger.warning(f"队列已满，任务入队失败: task_id={task_id}, queue_key={queue_key}, size={queue_size}")
                return False

            # 3. 存储任务元信息（任务详情、入队时间、状态）
            task_meta = {
                "task_id": task_id,
                "task_type": task["task_type"],
                "query": task["query"],
                "metadata": metadata,
                "enqueue_time": datetime.now().strftime(DEFAULT_DATETIME_FORMAT),
                "status": TASK_STATUSES.PENDING,  # 初始状态：待执行
                "retries": 0,  # 重试次数
                "timeout": metadata.get("timeout", self.task_timeout)
            }
            self.redis_client.hset(
                key=f"{self.task_meta_prefix}{task_id}",
                mapping=task_meta
            )
            # 设置元信息过期时间（任务超时后自动清理）
            self.redis_client.expire(
                key=f"{self.task_meta_prefix}{task_id}",
                time=task_meta["timeout"] + 3600  # 超时后额外保留1小时
            )

            # 4. 任务入队（Redis List：左进右出）
            self.redis_client.lpush(queue_key, task_id)
            logger.debug(f"任务入队成功: task_id={task_id}, priority={priority.value}, queue_key={queue_key}")

            # 5. 更新任务状态为"待执行"
            self.update_task_status(task_id, TASK_STATUSES.PENDING)
            return True
        except ParamError as e:
            logger.error(f"任务格式错误，入队失败: task_id={task.get('task_id')}, error={str(e)}")
            return False
        except Exception as e:
            logger.error(f"任务入队失败: task_id={task.get('task_id')}, error={str(e)}", exc_info=True)
            return False

    def dequeue_task(self, priority_order: List[TaskPriority] = None) -> Optional[Dict[str, Any]]:
        """
        任务出队（按优先级顺序获取任务，高优先级先执行）
        Args:
            priority_order: 优先级顺序（默认：HIGH > MEDIUM > NORMAL > LOW）
        Returns:
            出队的任务字典，无任务返回None
        """
        # 默认优先级顺序
        if not priority_order:
            priority_order = [TaskPriority.HIGH, TaskPriority.MEDIUM, TaskPriority.NORMAL, TaskPriority.LOW]

        try:
            # 1. 按优先级遍历队列，获取第一个非空队列的任务
            for priority in priority_order:
                queue_key = self._get_queue_key(priority)
                task_id = self.redis_client.rpop(queue_key)  # 右出（FIFO）
                if task_id:
                    task_id = task_id.decode("utf-8")
                    # 2. 获取任务元信息
                    task_meta = self.redis_client.hgetall(f"{self.task_meta_prefix}{task_id}")
                    if not task_meta:
                        logger.warning(f"任务元信息不存在，跳过: task_id={task_id}")
                        continue

                    # 3. 转换元信息类型（Redis返回字符串，需转义）
                    task_meta = self._parse_task_meta(task_meta)

                    # 4. 检查任务是否超时
                    enqueue_time = datetime.strptime(task_meta["enqueue_time"], DEFAULT_DATETIME_FORMAT)
                    timeout = task_meta["timeout"]
                    if (datetime.now() - enqueue_time).total_seconds() > timeout:
                        logger.warning(
                            f"任务超时，标记为失败: task_id={task_id}, enqueue_time={task_meta['enqueue_time']}")
                        self.update_task_status(task_id, TASK_STATUSES.FAILED, "任务超时未执行")
                        continue

                    # 5. 更新任务状态为"运行中"
                    self.update_task_status(task_id, TASK_STATUSES.RUNNING)
                    logger.debug(f"任务出队成功: task_id={task_id}, priority={priority.value}")
                    return task_meta
            logger.debug("所有队列无待执行任务")
            return None
        except Exception as e:
            logger.error(f"任务出队失败: error={str(e)}", exc_info=True)
            return None
    
    async def smart_dequeue_and_distribute(self, priority_order: List[TaskPriority] = None) -> Optional[Dict[str, Any]]:
        """
        智能出队并分发任务（使用智能任务分发器）
        Args:
            priority_order: 优先级顺序
        Returns:
            任务分配结果
        """
        if not self.use_smart_distribution:
            # 如果未启用智能分发，返回None让系统使用传统方式处理
            return None
            
        try:
            # 1. 获取任务
            task = self.dequeue_task(priority_order)
            if not task:
                return None
                
            # 2. 使用智能任务分发器分配任务
            result = await smart_task_distributor.distribute_task(task)
            
            # 3. 更新任务状态
            if result['status'] == 'allocated':
                allocated_agents = [agent[0] for agent in result['allocated_agents']]
                logger.info(f"任务 {task['task_id']} 已智能分配给: {allocated_agents}")
                self.update_task_status(
                    task['task_id'],
                    TASK_STATUSES.RUNNING,
                    f"智能分配给 {len(allocated_agents)} 个智能体"
                )
            elif result['status'] == 'error':
                logger.error(f"任务 {task['task_id']} 智能分配失败: {result.get('message')}")
                # 分配失败，回退到传统方式
                return task
                
            return result
            
        except Exception as e:
            logger.error(f"智能任务分发失败: {str(e)}", exc_info=True)
            # 出错时回退到传统方式
            return None

    def _parse_task_meta(self, raw_meta: Dict[bytes, bytes]) -> Dict[str, Any]:
        """解析Redis存储的任务元信息（字节转字符串/字典）"""
        parsed_meta = {}
        for k, v in raw_meta.items():
            key = k.decode("utf-8")
            value = v.decode("utf-8")

            # 特殊字段转换（metadata为字典字符串）
            if key == "metadata":
                try:
                    parsed_meta[key] = eval(value)  # 简化：实际用json.loads
                except:
                    parsed_meta[key] = {}
            elif key == "retries":
                parsed_meta[key] = int(value)
            elif key == "timeout":
                parsed_meta[key] = int(value)
            else:
                parsed_meta[key] = value
        return parsed_meta

    def update_task_status(self, task_id: str, status: str, message: str = "") -> bool:
        """更新任务状态"""
        try:
            task_key = f"{self.task_meta_prefix}{task_id}"
            if not self.redis_client.exists(task_key):
                logger.warning(f"任务不存在，无法更新状态: task_id={task_id}")
                return False

            # 更新状态和消息
            self.redis_client.hset(task_key, mapping={
                "status": status,
                "update_time": datetime.now().strftime(DEFAULT_DATETIME_FORMAT),
                "message": message
            })
            logger.debug(f"任务状态更新: task_id={task_id}, status={status}, message={message[:50]}")
            return True
        except Exception as e:
            logger.error(f"任务状态更新失败: task_id={task_id}, status={status}, error={str(e)}", exc_info=True)
            return False

    def retry_task(self, task_id: str) -> bool:
        """任务重试（失败后重新入队）"""
        try:
            task_key = f"{self.task_meta_prefix}{task_id}"
            task_meta = self.redis_client.hgetall(task_key)
            if not task_meta:
                logger.warning(f"任务不存在，无法重试: task_id={task_id}")
                return False

            task_meta = self._parse_task_meta(task_meta)
            current_retries = task_meta["retries"]
            max_retries = task_meta.get("metadata", {}).get("max_retries", 3)  # 默认最大3次重试

            if current_retries >= max_retries:
                logger.warning(f"重试次数耗尽，无法重试: task_id={task_id}, retries={current_retries}/{max_retries}")
                self.update_task_status(task_id, TASK_STATUSES.FAILED, "重试次数耗尽")
                return False

            # 更新重试次数
            self.redis_client.hincrby(task_key, "retries", 1)
            # 重新入队
            priority = TaskPriority(task_meta.get("metadata", {}).get("priority", "normal"))
            queue_key = self._get_queue_key(priority)
            self.redis_client.lpush(queue_key, task_id)
            # 更新状态为"待重试"
            self.update_task_status(task_id, TASK_STATUSES.RETRYING, f"第{current_retries + 1}次重试")

            logger.debug(f"任务重试成功: task_id={task_id}, retries={current_retries + 1}/{max_retries}")
            return True
        except Exception as e:
            logger.error(f"任务重试失败: task_id={task_id}, error={str(e)}", exc_info=True)
            return False

    def get_task_status(self, task_id: str) -> Optional[Dict[str, Any]]:
        """查询任务状态"""
        try:
            task_key = f"{self.task_meta_prefix}{task_id}"
            task_meta = self.redis_client.hgetall(task_key)
            if not task_meta:
                logger.warning(f"任务不存在: task_id={task_id}")
                return None
            return self._parse_task_meta(task_meta)
        except Exception as e:
            logger.error(f"查询任务状态失败: task_id={task_id}, error={str(e)}", exc_info=True)
            return None

    def clean_expired_tasks(self) -> int:
        """清理过期任务（超时且状态为失败/完成的任务）"""
        try:
            cleaned_count = 0
            # 遍历所有任务元信息键（模糊匹配）
            task_keys = self.redis_client.keys(f"{self.task_meta_prefix}*")
            for key in task_keys:
                task_id = key.decode("utf-8").replace(self.task_meta_prefix, "")
                task_meta = self._parse_task_meta(self.redis_client.hgetall(key))

                # 条件：状态为完成/失败 + 已过期
                status = task_meta.get("status")
                update_time = task_meta.get("update_time")
                if status in [TASK_STATUSES.COMPLETED, TASK_STATUSES.FAILED] and update_time:
                    update_dt = datetime.strptime(update_time, DEFAULT_DATETIME_FORMAT)
                    if (datetime.now() - update_dt).total_seconds() > 86400:  # 完成/失败后保留24小时
                        self.redis_client.delete(key)
                        cleaned_count += 1
                        logger.debug(f"清理过期任务: task_id={task_id}")
            logger.info(f"清理过期任务完成，共清理{cleaned_count}个")
            return cleaned_count
        except Exception as e:
            logger.error(f"清理过期任务失败: error={str(e)}", exc_info=True)
            return 0


# 单例实例（全局复用）
task_scheduler_agent = TaskSchedulerAgent()

# 异步初始化智能任务分发器
import asyncio
async def initialize_smart_distributor():
    """异步初始化智能任务分发器"""
    try:
        await smart_task_distributor.initialize()
        # 启动任务分发循环
        asyncio.create_task(smart_task_distributor.run_distribution_loop())
        logger.info("智能任务分发器初始化并启动成功")
    except Exception as e:
        logger.error(f"智能任务分发器初始化失败: {str(e)}", exc_info=True)

# 尝试异步初始化（在事件循环中执行）
try:
    import sys
    if sys.version_info >= (3, 7):
        # 检查是否已经在事件循环中
        if hasattr(asyncio, 'get_running_loop'):
            try:
                loop = asyncio.get_running_loop()
                if loop.is_running():
                    # 如果已有运行中的事件循环，将初始化任务提交到该循环
                    loop.create_task(initialize_smart_distributor())
                else:
                    # 如果事件循环存在但未运行，创建新任务
                    loop.run_until_complete(initialize_smart_distributor())
            except RuntimeError:
                # 没有运行中的事件循环，创建新的
                asyncio.run(initialize_smart_distributor())
        else:
            # 兼容旧版Python
            asyncio.run(initialize_smart_distributor())
    else:
        # 旧版Python可能需要特殊处理
        logger.warning("当前Python版本可能不完全支持异步初始化")
except Exception as e:
    logger.error(f"智能任务分发器初始化尝试失败: {str(e)}")
    # 失败时禁用智能分发，使用传统方式
    task_scheduler_agent.use_smart_distribution = False

# 测试入口
if __name__ == "__main__":
    test_logger = Logger.update_context(task_id="task_scheduler_test", agent_name="task_scheduler_agent")
    test_logger.info("开始测试任务调度Agent")

    # 测试任务
    test_task = {
        "task_id": "test_task_001",
        "task_type": "text_qa",
        "query": "分布式多Agent系统的核心优势是什么？",
        "metadata": {"priority": "high", "max_retries": 2}
    }

    try:
        # 1. 任务入队
        enqueue_ok = task_scheduler_agent.enqueue_task(test_task)
        test_logger.info(f"任务入队结果: {enqueue_ok}")

        # 2. 查询任务状态
        task_status = task_scheduler_agent.get_task_status("test_task_001")
        test_logger.info(f"任务状态（入队后）: {task_status['status']}")

        # 3. 任务出队
        dequeued_task = task_scheduler_agent.dequeue_task()
        test_logger.info(f"出队任务: {dequeued_task['task_id']}, 状态: {dequeued_task['status']}")

        # 4. 更新任务状态为完成
        update_ok = task_scheduler_agent.update_task_status("test_task_001", TASK_STATUSES.COMPLETED, "执行成功")
        test_logger.info(f"状态更新结果: {update_ok}")

        # 5. 清理过期任务
        cleaned_count = task_scheduler_agent.clean_expired_tasks()
        test_logger.info(f"清理过期任务数: {cleaned_count}")

        test_logger.success("任务调度Agent测试完成")
    except Exception as e:
        test_logger.error(f"测试失败: {e}")